#!/usr/bin/env python3
import argparse
import os
import sys
from time import sleep

import grpc

# Import P4Runtime lib from parent utils dir
# Probably there's a better way of doing this.
sys.path.append(
    os.path.join(os.path.dirname(os.path.abspath(__file__)),
                 '../../utils/'))
import p4runtime_lib.bmv2
import p4runtime_lib.helper
from p4runtime_lib.error_utils import printGrpcError
from p4runtime_lib.switch import ShutdownAllSwitchConnections

# 新增ecmp_group流规则
def writeGroupRules(p4info_helper, forward_sw, dst_ip_addr, acmp_base, ecmp_count):
    table_entry = p4info_helper.buildTableEntry(
        table_name="MyIngress.ecmp_group",
        match_fields={
            "hdr.ipv4.dstAddr": (dst_ip_addr, 32)
        },
        action_name="MyIngress.set_ecmp_select",
        action_params={
            "ecmp_base": ecmp_base,
            "ecmp_count": ecmp_count
        }
    )
    try:
        forward_sw.WriteTableEntry(table_entry)
        print("Installed ingress hash_range rule on %s" % sw.name)
    except grpc.RpcError as e:
        print("Error installing ingress hash_range rule on {}: {}".format(sw.name, e))

# 新增ecmp_nhop流规则
def writeNhopRules(p4info_helper, sw, ecmp_select, nhop_dmac, nhop_ipv4, port):
    table_entry = p4info_helper.buildTableEntry(
        table_name="MyIngress.ecmp_nhop",
        match_fields={
            "meta.ecmp_select": ecmp_select
        },
        action_name="MyIngress.set_nhop",
        action_params={
            "nhop_dmac": nhop_dmac,
            "nhop_ipv4": nhop_ipv4,
            "port": port
        }
    )
    try:
        sw.WriteTableEntry(table_entry)
        print("Installed ingress next_hop rule on %s" % sw.name)
    except grpc.RpcError as e:
        print("Error installing ingress next_hop rule on {}: {}".format(sw.name, e))


def writeDMAC(p4info_helper, sw, match_fields, action_params):
    table_entry = p4info_helper.buildTableEntry(
        table_name="MyEgress.send_frame",
        match_fields={
            "standard_metadata.egress_port": match_fields
        },
        action_name="MyEgress.rewrite_mac",
        action_params=action_params
    )
    try:
        sw.WriteTableEntry(table_entry)
        print("Installed egress dmac rule on %s" % sw.name)
    except grpc.RpcError as e:
        print("Error installing egress dmac rule on {}: {}".format(sw.name, e))

# 从交换机上的所有表中读取表条目
def readTableRules(p4info_helper, sw):
    print('\n----- Reading tables rules for %s -----' % sw.name)
    for response in sw.ReadTableEntries():
        for entity in response.entities:
            entry = entity.table_entry
            table_name = p4info_helper.get_tables_name(entry.table_id)
            print('%s: ' % table_name, end=' ')

            for m in entry.match:
                print(p4info_helper.get_match_field_name(table_name, m.field_id), end=' ')
                print('%r' % (p4info_helper.get_match_field_value(m),), end=' ')
                action = entry.action.action
                action_name = p4info_helper.get_actions_name(action.action_id)
                print('->', action_name, end=' ')

            for p in action.params:
                print(p4info_helper.get_action_param_name(action_name, p.param_id), end=' ')
                print('%r' % p.value, end=' ')
                print()


def main(p4info_file_path, bmv2_file_path):
    p4info_helper = p4runtime_lib.helper.P4InfoHelper(p4info_file_path)

    try:
        # Create a switch connection object for each switch
        switches = []
        switch_conns = []

        for i in range(1, 7):
            sw = p4runtime_lib.bmv2.Bmv2SwitchConnection(
                name='s{}'.format(i),
                address='127.0.0.1:5005{}'.format(i),
                device_id=i - 1,
                proto_dump_file='logs/s{}-p4runtime-requests.txt'.format(i)
            )
            switches.append(sw)

            # 将每个交换机连接对象添加到列表中
            switch_conns.append(sw)

        for switch_conn in switch_conns:
            # 发送主控制权更新消息以将此控制器建立为主控制器（在执行任何其他写操作之前，P4Runtime要求执行此操作）
            switch_conn.MasterArbitrationUpdate()

        # 在交换机上安装P4程序
            switch_conn.SetForwardingPipelineConfig(
                p4info=p4info_helper.p4info,
                bmv2_json_file_path=bmv2_file_path
            )

            print(f"Installed P4 Program using SetForwardingPipelineConfig on {switch_conn}")

        # 新增ecmp_groups流表
        ecmp_groups = [
            (switch_conns[0], "10.0.0.1", 2, 6),  # h1发
            (switch_conns[0], "10.0.2.2", 2, 6),
            (switch_conns[0], "10.0.1.1", 1, 2),  # h1收
            (switch_conns[1], "10.0.1.1", 0, 1),
            (switch_conns[1], "10.0.2.2", 1, 2),
            (switch_conns[2], "10.0.1.1", 0, 1),
            (switch_conns[2], "10.0.2.2", 1, 2),
            (switch_conns[3], "10.0.1.1", 0, 1),
            (switch_conns[3], "10.0.2.2", 1, 2),
            (switch_conns[4], "10.0.1.1", 0, 1),
            (switch_conns[4], "10.0.2.2", 1, 2),
            (switch_conns[5], "10.0.2.2", 1, 2),  # h2收
            (switch_conns[5], "10.0.0.2", 2, 6),  # h2发
            (switch_conns[5], "10.0.1.1", 2, 6)
        ]
        for forward_sw, dst_ip_addr, ecmp_base, ecmp_count in ecmp_groups:
            writeGroupRules(forward_sw=forward_sw,
                            dst_ip_addr=dst_ip_addr, ecmp_base=ecmp_base,
                            ecmp_count=ecmp_count,
                            p4info_helper=p4info_helper)
        ecmp_nhops = [
            (switch_conns[0], 1, "08:00:00:00:01:01","10.0.1.1", 1),
            (switch_conns[0], 2, "08:00:00:00:02:00","10.0.2.2", 2),
            (switch_conns[0], 3, "08:00:00:00:02:00","10.0.2.2", 3),
            (switch_conns[0], 4, "08:00:00:00:02:00","10.0.2.2", 4),
            (switch_conns[0], 5, "08:00:00:00:02:00","10.0.2.2", 5),
            (switch_conns[1], 0, "08:00:00:00:01:00","10.0.1.1", 1),  # s2 to s1
            (switch_conns[1], 1, "08:00:00:00:02:00","10.0.2.2", 2),  # s2 to s6
            (switch_conns[2], 0, "08:00:00:00:01:00","10.0.1.1", 1),  # s3 to s1
            (switch_conns[2], 1, "08:00:00:00:02:00","10.0.2.2", 2),  # s3 to s6
            (switch_conns[3], 0, "08:00:00:00:01:00","10.0.1.1", 1),  # s4 to s1
            (switch_conns[3], 1, "08:00:00:00:02:00","10.0.2.2", 2),  # s4 to s6
            (switch_conns[4], 0, "08:00:00:00:01:00","10.0.1.1", 1),  # s5 to s1
            (switch_conns[4], 1, "08:00:00:00:02:00","10.0.2.2", 2),  # s5 to s6
            (switch_conns[5], 1, "08:00:00:00:02:01","10.0.2.2", 1),
            (switch_conns[5], 2, "08:00:00:00:01:00","10.0.1.1", 2),
            (switch_conns[5], 3, "08:00:00:00:01:00","10.0.1.1", 3),
            (switch_conns[5], 4, "08:00:00:00:01:00","10.0.1.1", 4),
            (switch_conns[5], 5, "08:00:00:00:01:00","10.0.1.1", 5)
        ]
        for forward_sw, ecmp_select, nhop_dmac, nhop_ipv4, port in ecmp_nhops:
            writeNhopRules(p4info_helper=p4info_helper,
                           forward_sw=forward_sw,
                           ecmp_select=ecmp_select,
                           nhop_dmac=nhop_dmac,
                           nhop_ipv4=nhop_ipv4, port=port)
        for switch_conn in switch_conns:
            readTableRules(p4info_helper=p4info_helper,sw=switch_conn)

    except KeyboardInterrupt:
        print(" Shutting down.")
    except grpc.RpcError as e:
        printGrpcError(e)

    ShutdownAllSwitchConnections()

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='P4Runtime Controller')
    parser.add_argument('--p4info', help='p4info proto in text format from p4c',
                        type=str, action="store", required=False,
                        default='./build/load_balance.p4.p4info.txt')
    parser.add_argument('--bmv2-json', help='BMv2 JSON file from p4c',
                        type=str, action="store", required=False,
                        default='./build/load_balance.json')
    args = parser.parse_args()

    if not os.path.exists(args.p4info):
        parser.print_help()
        print("\np4info file not found: %s\nHave you run 'make'?" % args.p4info)
        parser.exit(1)
    if not os.path.exists(args.bmv2_json):
        parser.print_help()
        print("\nBMv2 JSON file not found: %s\nHave you run 'make'?" % args.bmv2_json)
        parser.exit(1)
    main(args.p4info, args.bmv2_json)